import cv2
import numpy as np
from utils.common import load_image, show_image, make_dirs
from utils.config import exp5_output_dir, exp5_light_path

# 确保输出目录存在
make_dirs(exp5_output_dir)

# 加载图片
image = load_image()

# 将图像转换为灰度图
gray_image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)

# 创建随机颜色
color = np.random.randint(0, 255, (100, 3))

# ShiTomasi角点检测参数
feature_params = dict(maxCorners=100,
                      qualityLevel=0.3,
                      minDistance=7,
                      blockSize=7)

# Lucas-Kanade光流参数
lk_params = dict(winSize=(15, 15),
                 maxLevel=2,
                 criteria=(cv2.TERM_CRITERIA_EPS | cv2.TERM_CRITERIA_COUNT, 10, 0.03))

# 检测角点
p0 = cv2.goodFeaturesToTrack(gray_image, mask=None, **feature_params)

# 创建一个掩码图像用于绘图
mask = np.zeros_like(image)

# 假设我们有一个连续的帧，这里简单使用同一帧代替
next_gray = gray_image.copy()

# 计算光流
p1, st, err = cv2.calcOpticalFlowPyrLK(gray_image, next_gray, p0, None, **lk_params)

# 选择好的点
good_new = p1[st == 1]
good_old = p0[st == 1]

# 绘制轨迹
for i, (new, old) in enumerate(zip(good_new, good_old)):
    a, b = new.ravel()
    c, d = old.ravel()
    mask = cv2.line(mask, (int(a), int(b)), (int(c), int(d)), color[i].tolist(), 2)
    image = cv2.circle(image, (int(a), int(b)), 5, color[i].tolist(), -1)

img = cv2.add(image, mask)

show_image("Optical Flow Feature Image", img)

# 保存结果
cv2.imwrite(exp5_light_path, img)
print(f"基于光流的特征提取结果已保存到：{exp5_light_path}")